import pandas as pd
import numpy as np

def convert_new_format(df: pd.DataFrame) -> pd.DataFrame:
    '''转换新格式的Excel数据为标准格式'''
    try:
        df = df.copy()
        
        # 第一列是用电量
        df.iloc[:, 0] = pd.to_numeric(df.iloc[:, 0].astype(str).str.replace(',', ''))
        # 第二列日期
        df.iloc[:, 1] = pd.to_datetime(df.iloc[:, 1]).dt.strftime('%Y-%m')
        # 交换列的位置
        df = df.iloc[:, [1, 0]]
        df.columns = range(len(df.columns))
        return df    
    except Exception as e:
        print(f"格式转换错误: {str(e)}")
        return None

def read_data(file_path: str, is_new_format: bool = False) -> pd.DataFrame:
    '''读取excel数据
    
    Args:
        file_path: Excel文件路径
        is_new_format: 是否为新格式的Excel文件
        
    Returns:
        DataFrame: 第一列为年月(YYYY-MM格式)，第二列为用电量
    '''
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path)
        df = df.iloc[:, :2]
        
        if df.empty:
            raise ValueError("Excel文件中没有数据") 
        if is_new_format:
            df = convert_new_format(df)
            if df is None:
                raise ValueError("格式转换失败")
        else:
            # 日期列为字符串格式
            df.iloc[:, 0] = pd.to_datetime(df.iloc[:, 0]).dt.strftime('%Y-%m')
            df.iloc[:, 1] = pd.to_numeric(df.iloc[:, 1])

        # 验证
        if df.iloc[:, 1].isna().any():
            raise ValueError("存在无效的用电量数据")
        if df.iloc[:, 0].isna().any():
            raise ValueError("存在无效的日期数据")
        return df       
    except Exception as e:
        print(f"读取Excel文件时发生错误: {str(e)}")
        return None